package com.yystudy.aiagentbackend.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import lombok.extern.slf4j.Slf4j;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author: yp123
 * @since: 2025/8/6
 * 基于本地文件存储的对话记忆
 */
@Slf4j
public class FileBasedChatMemory implements ChatMemory {

    /**
     * 基础目录
     */
    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    /**
     * 构造对象时，指定文件保存目录
     * @param dir
     */
    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()){
            baseDir.mkdirs();
        }
    }

    /**
     * 存储添加消息
     * @param conversationId  会话id
     * @param messages   消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    /**
     * 根据会话id获取消息
     * @param conversationId  会话id
     */
    @Override
    public List<Message> get(String conversationId) {
        return getOrCreateConversation(conversationId);
    }

    /**
     * 删除会话
     * @param conversationId  会话id
     */
    @Override
    public void clear(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        if (conversationFile.exists()){
            conversationFile.delete();
        }
    }

    private List<Message> getOrCreateConversation(String conversationId){
        File conversationFile = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (conversationFile.exists()){
            try (Input input = new Input(new FileInputStream(conversationFile))){
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e){
                log.error("error：{}", e.getMessage(), e);
            }
        }
        return messages;
    }

    private void saveConversation(String conversationId, List<Message> messages){
        File conversationFile = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(conversationFile))){
            kryo.writeObject(output, messages);
        } catch (IOException e){
            log.error("error：{}", e.getMessage(), e);
        }
    }


    private File getConversationFile(String conversationId){
        return new File(BASE_DIR, conversationId + ".kryo");
    }


}
